import argparse
import consul


class ConsulKV:
    @classmethod
    def __get_ops(cls):
        args_parse = argparse.ArgumentParser(description="consul连接信息")
        args_parse.add_argument(
            '--consul_host', type=str, required=False, default="127.0.0.1", help="consul连接地址, 默认本地"
        )
        args_parse.add_argument(
            '--consul_port', type=int, required=False, default=8500, help="consul连接端口，默认8500"
        )
        args_parse.add_argument(
            '--consul_token', type=str, required=False, default=None, help="consul认证token，默认为None"
        )
        args_parse.add_argument(
            '--consul_dc', type=str, required=False, default="dc1", help="consul数据中心，默认为dc1"
        )
        return args_parse.parse_args()

    def __init__(self):
        args = self.__get_ops()
        self._conn = consul.Consul(
            host=args.consul_host,
            port=args.consul_port,
            token=args.consul_token,
            dc=args.consul_dc
        )

    def get(self, key):
        _, keys = self._conn.kv.get(key=key, recurse=True)
        return {d['Key']: d['Value'].decode('utf-8') for d in keys if d['Value'] is not None}

    def put(self, key, value):
        return self._conn.kv.put(key, value)

    def delete(self, key, recurse=False):
        return self._conn.kv.delete(key, recurse)
